# This is to count the size of the parameters of the model
import os,sys
root_path = os.getcwd()
current_path = os.path.join(root_path,'finetune_src')
sys.path.append(root_path)
sys.path.append(current_path)

import torch
import torch.nn as nn
import numpy as np
from r2r_geo_slot.parser import parse_args
from models.model_HAMT_GEO import VLNBertCMT, Critic
from models.vilmodel_cmt import NavCMT
from r2r_geo_slot.slot_attention import SlotAttention
from einops import rearrange, repeat
from transformers import PretrainedConfig
from thop import profile
import time

if __name__ == '__main__':
    args = parse_args()
    ''' Parameter count '''
    vln_bert = VLNBertCMT(args).cuda()
    critic = Critic(args).cuda()
    slot_attn = SlotAttention(
        num_slots=1,
        dim=(768+4)
    ).cuda()
    vln_bert_size = sum([x.numel() for x in vln_bert.parameters()])
    critic_size = sum([x.numel() for x in critic.parameters()])
    slot_attn_size = sum([x.numel() for x in slot_attn.parameters()])
    print('Parameters:')
    print(vln_bert_size)
    print(critic_size)
    print(slot_attn_size)
    print('Total size %.2f' %(vln_bert_size+critic_size+slot_attn_size))

    '''FLOPs count for specific proposed module'''
    # propose_model = ProposedModule(args)
    # txt_embeds = torch.rand(8,60,768)
    # act_embeds = torch.rand(8,37,768)
    # obj_embeds = torch.rand(8,44,768)

    # obj_img_fts = torch.rand(8,36,3,8)
    # nav_types = torch.ones(8,36,dtype=torch.int)

    # traj_ori_embeds = torch.rand(8,36,768)
    
    # gmap_embeds = torch.rand(8,36,768)
    # vp_embeds = torch.rand(8,36,768)
    # txt_embeds = torch.rand(8,60,768)

    # # obj_img_fts, nav_types, obj_lens
    # macs, params = profile(propose_model, (gmap_embeds,vp_embeds,txt_embeds))
    # print('GFLOPs:%.3f' %(macs*2/(10**9)))
    # print(1)

    '''FLOPs count for the whole proposed module '''

    '''Language'''
    forward_time = 0
    start_time = time.time()
    txt_ids = torch.randint(0,2000,size=(8,44)).cuda()
    txt_masks = torch.randint(0, 2, size=(8, 44), dtype=torch.bool).cuda()
    mode = 'language'
    lan_input = {
        'txt_ids': txt_ids,
        'txt_masks': txt_masks,
    }
    forward_time += (time.time()-start_time) 
    lan_macs, lan_params = profile(vln_bert, (mode, txt_ids, txt_masks))
    lan_gflops = lan_macs*2/(10**9)
    print('Language GFLOPs:%.3f' %(lan_gflops))
    print('*****')

    '''history'''
    start_time = time.time()
    hist_img_feats = torch.rand(8,2304).cuda()
    hist_ang_feats = torch.rand(8,4).cuda()
    ob_step = 4
    hist_pano_img_feats = torch.rand(8,36,2304).cuda()
    hist_pano_ang_feats = torch.rand(8,36,4).cuda()
    mode = 'history'

    history_input = {
        'hist_img_feats': hist_img_feats,
        'hist_ang_feats': hist_ang_feats,
        'ob_step': ob_step,
        'hist_pano_img_feats': hist_pano_img_feats,
        'hist_pano_ang_feats': hist_pano_ang_feats,
        'mode': mode
    }

    pan_macs, pan_params = profile(vln_bert,(mode, None, None, None, hist_img_feats, hist_ang_feats, hist_pano_img_feats, hist_pano_ang_feats,None,None,ob_step))
    forward_time += (time.time()-start_time) 
    pan_gflops = pan_macs * 2 / (10 ** 9)
    print('Panorama GFLOPs: %.3f' % (pan_gflops))
    print('*****')

    '''visual'''
    start_time = time.time()
    txt_embeds = torch.rand(8,44,768).cuda()
    txt_masks = torch.rand(8,44).cuda()
    hist_embeds = [torch.rand(8,768).cuda() for _ in range(6)]
    hist_masks = [torch.randint(0, 2, size=(8, 1), dtype=torch.bool).cuda() for _ in range(6)]
    ob_img_feats = torch.rand(8,39,2304).cuda()
    ob_ang_feats = torch.rand(8,39,4).cuda()
    ob_nav_types = torch.randint(0,3,size=(8,39)).cuda()
    ob_masks = torch.randint(0, 2, size=(8, 39), dtype=torch.bool).cuda()
    hist_lens = [6 for _ in range(8)]
    mode = 'visual'

    visual_input = {
        'txt_embeds': txt_embeds,
        'txt_masks': txt_masks,
        'hist_embeds': hist_embeds,
        'hist_masks': hist_masks,
        'ob_img_feats': ob_img_feats,
        'ob_ang_feats': ob_ang_feats,
        'ob_nav_types': ob_nav_types,
        'ob_masks': ob_masks,
        'mode': mode
    }

    nav_macs, nav_params = profile(vln_bert, (mode, txt_ids, txt_masks, txt_embeds, hist_img_feats, hist_ang_feats, hist_pano_img_feats, hist_pano_ang_feats, hist_embeds, hist_lens,ob_step,ob_img_feats,ob_ang_feats,ob_nav_types,ob_masks))
    forward_time += (time.time()-start_time) 
    nav_gflops = nav_macs * 2 / (10 ** 9)
    print('Panorama GFLOPs: %.3f' % (nav_gflops))
    print('*****')

    '''Slot Attn'''
    cand_feat = torch.rand(8, 39, 768+4).cuda()
    pano_feat = torch.rand(8, 36, 768+4).cuda()
    slot_macs, slot_params = profile(slot_attn, (cand_feat, pano_feat))
    forward_time += (time.time()-start_time) 
    slot_gflops = slot_macs * 2 / (10 ** 9)
    print('Panorama GFLOPs: %.3f' % (slot_gflops))
    print('*****')
    
    print('Total GFlops: %.3f' % (lan_gflops + pan_gflops + nav_gflops + slot_gflops))
    print('Forward Time: %.3f s' % forward_time) 


